from math import acosh

import yaml
import argparse
import os
import re

script_path = os.path.dirname(os.path.abspath(__file__))


def verbose(message: str):
    if args.verbose:
        print(message)


class SmokeCaseConfigs:
    def __init__(self, inventory_path, template_path, cli_path):
        self.inventory_path = inventory_path
        self.template_path = template_path
        self.cli_path = cli_path
        self.configs = list()
        self.docker_containers = dict()
        self.common_configs = dict()
        self.node_log_path = "${LOG_PATH}"

    def insert_hosts(self):
        try:
            if self.inventory_path:
                with open(self.inventory_path, 'r', encoding='utf-8') as file:
                    data = yaml.safe_load(file)
                    _vars = data['all']['vars']
                    for _var, _var_value in _vars.items():
                        if _var == 'ansible_ssh_private_key_file':
                            self.common_configs.update({'ssh_private_key': _var_value})
                        if _var == 'proxy_port':
                            self.common_configs.update({'PROXY_PORT': _var_value})
                    children = data['all']['children']
                    for key, child in children.items():
                        if child.get('hosts'):
                            self.node_config_update(child['hosts'])
                        elif child.get('children'):
                            sub_children = child.get('children')
                            for sub_key, sub_child in sub_children.items():
                                if sub_child.get('hosts'):
                                    self.node_config_update(sub_child['hosts'])
                        else:
                            # todo 遍历所有子节点，直到找到每一个hosts
                            continue
            elif self.cli_path:
                with open(self.cli_path, 'r', encoding='utf-8') as file:
                    data = yaml.safe_load(file)
                    children = data['all']['children']
                    self.cli_node_config_update(children['C']['hosts'])
                    self.cli_node_config_update(children['P']['hosts'])
                    self.cli_node_config_update(children['D']['hosts'])
            else:
                pass

        except yaml.YAMLError as e:
            verbose(f"inventory 解析错误: {e}")
            exit(-1)
        except FileNotFoundError:
            verbose("inventory文件未找到")
            exit(-1)

    def cli_node_config_update(self, child_hosts: dict):
        for sub_child, sub_value in child_hosts.items():
            ansible_host = sub_value.get('ansible_host')
            node_config = {
                "name": sub_child.upper(),
                "host": ansible_host,
                "container_suffix": sub_child.lower(),
                "container_name": sub_value.get('container_name'),
            }
            self.common_configs.update({'ssh_private_key': sub_value.get('ansible_ssh_private_key_file')})
            self.configs.append(node_config)

    def node_config_update(self, child_hosts: dict):
        for sub_child, sub_value in child_hosts.items():
            ansible_host = sub_value.get('ansible_host')
            node_config = {
                "name": sub_child.upper(),
                "host": ansible_host,
                "container_suffix": sub_child.lower()
            }
            self.configs.append(node_config)

    def insert_containers(self):
        try:
            with open(self.template_path, 'r', encoding='utf-8') as file:
                data = yaml.safe_load(file)
                # 取container名称
                template_envs = data[0]['environment']
                p_container_name = template_envs["DOCKER_NAME_P"]
                d_container_name = template_envs["DOCKER_NAME_D"]
                c_container_name = template_envs["DOCKER_NAME_C"]
                for item in self.configs:
                    if 'P' in item['name']:
                        item['container_name'] = p_container_name + "_" + item['container_suffix']
                    elif 'D' in item['name']:
                        item['container_name'] = d_container_name + "_" + item['container_suffix']
                    elif 'C' in item['name']:
                        item['container_name'] = c_container_name + "_" + item['container_suffix']
                    else:
                        verbose(f"no match pd name, item: {item}")
                self.docker_containers.update({
                    "p_container_name": p_container_name,
                    "d_container_name": d_container_name,
                    "c_container_name": c_container_name,
                })
                # 取log输出路径结构
                vars = data[0]['vars']
                run_vllm_server_prefill_cmd = vars['run_vllm_server_prefill_cmd']
                self.node_log_path = self.get_log_path(run_vllm_server_prefill_cmd)
        except yaml.YAMLError as e:
            verbose(f"inventory 解析错误: {e}")
            exit(-1)
        except FileNotFoundError:
            verbose("inventory文件未找到")
            exit(-1)

    @staticmethod
    def get_log_path(cmd_str: str):
        node_log_path = "${LOG_PATH}"
        cmd_list = cmd_str.split('\n')
        for line in cmd_list:
            if "${LOG_PATH}" in line and "run_prefill.log" in line:
                pattern = r"> (.*?)/run_prefill\.log"
                match = re.search(pattern, line)
                if match:
                    node_log_path = match.group(1)
                    verbose(f"匹配node log path结果: {node_log_path}")
                else:
                    verbose("未找到匹配内容")
            continue
        return node_log_path

    def parse_configs(self):
        verbose(f"接收inventory文件: {self.inventory_path}, 接收server_template文件: {self.template_path}")
        verbose("开始解析inventory文件....")
        self.insert_hosts()
        verbose("inventory文件解析完成")
        if os.path.exists(args.template_path):
            verbose("开始解析server_template文件....")
            self.insert_containers()
            verbose("server_template文件解析完成")
        verbose(f"解析完成，configs: {self.configs}")

    def trans_to_export_cmd(self):
        export_content = ""
        for item in self.configs:
            export_content += f"export {item['name']}=\"{item['host']}\"\n"
            if item.get('container_name'):  # cli分支存在container_name，inventory分支不存在
                export_content += f"export {item['container_suffix']}_container=\"{item['container_name']}\"\n"
        for key, value in self.docker_containers.items():
            export_content += f"export {key}=\"{value}\"\n"
        for key, value in self.common_configs.items():
            export_content += f"export {key}=\"{value}\"\n"
        export_content += f"export node_log_path=\"{self.node_log_path}\"\n"
        return export_content


def generate_export_cmd(inventory_path, template_path):
    smoke_case_configs = SmokeCaseConfigs(inventory_path, template_path, "")
    # 解析inventory和template
    smoke_case_configs.parse_configs()
    return smoke_case_configs.trans_to_export_cmd()

def generate_cli_export_cmd(cli_path: str):
    smoke_case_configs = SmokeCaseConfigs("", "", cli_path)
    # 解析inventory和template
    smoke_case_configs.parse_configs()
    return smoke_case_configs.trans_to_export_cmd()

if __name__ == '__main__':
    args = argparse.ArgumentParser()
    args.add_argument('-i', '--inventory_path', required=False, help='path to inventory.yml', default="")
    args.add_argument('-t', '--template_path', required=False, help='path to server_template.yml', default="")
    args.add_argument('-c', '--cli_path', required=False, help='path to cli config.yml', default="cli_ds_r1_w8a8c16_a2_1p1d_720tps.yaml")
    args.add_argument('--verbose', action='store_true', help="verbose mode")
    args = args.parse_args()

    # cli 用例解析
    if args.cli_path and os.path.exists(args.cli_path):
        export_content = generate_cli_export_cmd(args.cli_path)
        print(f"{export_content}")
    else:
        if not os.path.exists(args.inventory_path):
            print('Inventory file does not exist')
            exit(-1)

        if not os.path.exists(args.template_path):
            print('Template file does not exist')

        export_content = generate_export_cmd(args.inventory_path, args.template_path)
        print(f"{export_content}")